/* ========== User Space Loader (loader.c) ========== */
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <signal.h>
#include <sys/resource.h>
#include <bpf/libbpf.h>
#include <bpf/bpf.h>

static volatile bool exiting = false;

static void sig_handler(int sig)
{
    exiting = true;
}

static int libbpf_print_fn(enum libbpf_print_level level, const char *format, va_list args)
{
    return vfprintf(stderr, format, args);
}

int main(int argc, char **argv)
{
    struct bpf_object *obj = NULL;
    //struct bpf_program *prog;
    struct bpf_map *map;
    struct bpf_link *link = NULL;
    int err;

    /* 设置信号处理 */
    signal(SIGINT, sig_handler);
    signal(SIGTERM, sig_handler);

    /* 设置 libbpf 调试输出 */
    libbpf_set_print(libbpf_print_fn);

    /* 提升内存限制，避免加载失败 */
    struct rlimit rlim_new = {
        .rlim_cur = RLIM_INFINITY,
        .rlim_max = RLIM_INFINITY,
    };
    if (setrlimit(RLIMIT_MEMLOCK, &rlim_new)) {
        fprintf(stderr, "Failed to increase RLIMIT_MEMLOCK limit: %s\n", strerror(errno));
        return 1;
    }

    /* 打开 BPF 对象文件 */
    obj = bpf_object__open_file("bpf_handle_try_charge_memcg.bpf.o", NULL);
    if (libbpf_get_error(obj)) {
        fprintf(stderr, "Failed to open BPF object file: %s\n", strerror(errno));
        return 1;
    }

    /* 加载 BPF 程序到内核 */
    err = bpf_object__load(obj);
    if (err) {
        fprintf(stderr, "Failed to load BPF object: %s\n", strerror(-err));
        goto cleanup;
    }

    printf("BPF object loaded successfully\n");

    /* 查找 struct_ops map */
    map = bpf_object__find_map_by_name(obj, "try_charge_ops");
    if (!map) {
        fprintf(stderr, "Failed to find struct_ops map 'try_charge_ops'\n");
        err = -ENOENT;
        goto cleanup;
    }

    /* 创建 struct ops 链接 (自动注册) */
    link = bpf_map__attach_struct_ops(map);
    if (libbpf_get_error(link)) {
        fprintf(stderr, "Failed to attach struct ops: %s\n", strerror(-libbpf_get_error(link)));
        link = NULL;
        err = -1;
        goto cleanup;
    }

    printf("Struct ops attached successfully\n");
    printf("Press Ctrl+C to exit...\n");

    /* 保持程序运行 */
    while (!exiting) {
        sleep(1);
    }

    printf("\nDetaching struct ops...\n");

cleanup:
    /* 清理资源 */
    if (link)
        bpf_link__destroy(link);
    if (obj)
        bpf_object__close(obj);

    printf("Program exited\n");
    return err;
}
